"""
FEniCS tutorial demo program: Poisson equation with Dirichlet conditions.
Test problem is chosen to give an exact solution at all nodes of the mesh.

  -Laplace(u) = f    in the unit square
            u = u_D  on the boundary

  u_D = 1 + x^2 + 2y^2
    f = -6


ref: https://www.iesensor.com/blog/2017/05/24/gmsh_fenics_meshing/
     https://comphysblog.wordpress.com/2018/08/15/fenics-2d-electrostatics-with-imported-mesh-and-boundaries/
"""

from __future__ import print_function
from fenics import *
from dolfin import *
import matplotlib.pyplot as plt
import os
from vtkplotter.dolfin import plot


#fname = "t1"
fname = "hl2a"
mesh = Mesh(fname+".xml")
if os.path.exists( fname+"_physical_region.xml"):
    subdomains = MeshFunction("size_t", mesh, fname+"_physical_region.xml")
    #plot(subdomains)
if os.path.exists( fname+"_facet_region.xml"):
    boundaries = MeshFunction("size_t", mesh, fname+"_facet_region.xml")
    #plot(boundaries)

#plot(mesh)
#plt.show()


print("xml mesh reading done")


V = FunctionSpace(mesh, 'P', 1)



# Define boundary condition
u_D = Expression('1 + x[0]*x[0] + 2*x[1]*x[1]', degree=2)

def boundary(x, on_boundary):
    return on_boundary

#bc = DirichletBC(V, u_D, boundary)

#the parameter after boundaries, see hl2a_facet_region.xml file
inner_edge_boundary = DirichletBC(V, Constant(100.0), boundaries, 3)
target_boundary = DirichletBC(V, Constant(50.0), boundaries, 1)
#first_wall_boundary = DirichletBC(V, Constant(0.0), boundaries, 2)


#bcs =[inner_edge_boundary, target_boundary, first_wall_boundary]
bcs =[inner_edge_boundary, target_boundary]





# Define variational problem
u = TrialFunction(V)
v = TestFunction(V)
f = Constant(10.0)
a = dot(grad(u), grad(v))*dx
L = f*v*dx

# Compute solution
u = Function(V)
solve(a == L, u, bcs)

# Plot solution and mesh
#plot(u)
#plot(mesh)

#plot using vtk
plot(u, mode='color', style=1)


# Save solution to file in VTK format
vtkfile = File('poisson/solution.pvd')
vtkfile << u

# Compute error in L2 norm
error_L2 = errornorm(u_D, u, 'L2')

# Compute maximum error at vertices
vertex_values_u_D = u_D.compute_vertex_values(mesh)
vertex_values_u = u.compute_vertex_values(mesh)
import numpy as np
error_max = np.max(np.abs(vertex_values_u_D - vertex_values_u))

# Print errors
print('error_L2  =', error_L2)
print('error_max =', error_max)

# Hold plot
plt.show()
